import sys, os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
from tqdm import tqdm
import json
from preprossData.utils.crop_padding import crop_resize_save

def gather_data(data_path, tar_path, prefix='GlaucomaFundus_'):
    """
    处理 Glaucoma Fundus 数据集，将各类别图片预处理后保存，并生成标注信息 JSON 文件。
    
    数据集要求 data_path 下存在以下文件夹：
      - advanced_glaucoma
      - early_glaucoma
      - normal_control
      
    folder_mapped 定义了每个文件夹对应的诊断文本：
      - 'advanced_glaucoma' -> "advanced glaucoma"
      - 'early_glaucoma'    -> "early glaucoma"
      - 'normal_control'    -> "normal"
    
    每张图片经过 resize_save 处理后统一保存为 PNG 格式，新文件名格式为：{prefix}{原图基础名}.png
    
    最终生成的 annotations.json 中保存每个图片的标注信息。
    
    参数：
      data_path (str): 原始数据集根目录
      tar_path (str): 预处理后数据存放目录
      prefix (str): 新图片名称前缀，默认为 "GlaucomaFundus_"
    
    返回：
      dict: 标注信息字典
    """
    # 创建目标目录及 images 子目录
    os.makedirs(tar_path, exist_ok=True)
    images_dir = os.path.join(tar_path, "images")
    os.makedirs(images_dir, exist_ok=True)
    
    data_dict = {}
    
    folder_mapped = {
        'advanced_glaucoma': "advanced glaucoma",
        'early_glaucoma': "early glaucoma",
        'normal_control': "normal",
    }
    
    for folder_name, diagnosis_text in folder_mapped.items():
        image_path_root = os.path.join(data_path, folder_name)
        if not os.path.exists(image_path_root):
            print(f"Warning: folder '{image_path_root}' does not exist, skip.")
            continue
        # 获取该文件夹下所有图片文件名
        image_name_list = sorted(os.listdir(image_path_root))
        total_image_num=len(image_name_list)
        for image_name in tqdm(image_name_list,total=total_image_num,desc=f"{folder_name} process",unit='images'):
            # 处理文件名中可能包含多个点，确保正确获取基础名称
            base_name, _ = os.path.splitext(image_name)
            new_image_name = f"{prefix}{base_name}.png"
            src_image_path = os.path.join(image_path_root, image_name)
            dest_image_path = os.path.join(images_dir, new_image_name)
            
            # 对图像进行 resize 保存，resize_save 为预处理函数（统一输出为 PNG 格式）
            crop_info=crop_resize_save(
                image_path=src_image_path,
                save_path=dest_image_path,
                resize=(224, 224)
            )
            
            data_dict[new_image_name] = {
                'image_name': new_image_name,
                'image_path': os.path.join('images', new_image_name),
                'original_path':os.path.relpath(src_image_path,data_path),
                'crop_info':crop_info,
                'diagnosis': {'text': diagnosis_text}
            }
    
    # 保存标注信息到 JSON 文件
    annotations_path = os.path.join(tar_path, "annotations.json")
    with open(annotations_path, "w", encoding="utf-8") as f:
        json.dump(data_dict, f, indent=4)
    
    return data_dict

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Glaucoma Fundus 数据集预处理程序")
    parser.add_argument("--data_path", type=str, required=True,
                        help="原始数据集根目录（应包含 advanced_glaucoma, early_glaucoma, normal_control 文件夹）")
    parser.add_argument("--tar_path", type=str, required=True,
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="GlaucomaFundus_",
                        help="处理后图片名称前缀，默认为 'GlaucomaFundus_'")
    args = parser.parse_args()
    
    annotations = gather_data(args.data_path, args.tar_path, prefix=args.prefix)
    print("Preprocessing completed. Annotations saved.")
